
#include <vector>
#include <iostream>
#include <fstream>
#include <string>
#include <opencv2/opencv.hpp>
#include "opencv2/xfeatures2d/nonfree.hpp"

#define DEBUG

#define IMAGE_FILE(index) "image_" + std::to_string(index) + ".jpg"
#define POSE_FILE(index) "pose_" + std::to_string(index) + ".yaml"

#define CAPTURE_WIN "Capturing... Press ESC to stop"
#define MATCH_WIN "Match"
#define TOTAL_WIN "Total"
#define BEFORE_WARPING_WIN "Before warpping"
#define AFTER_WARPING_WIN "After warpping"
#define WEIGHT_WIN "Weight"

#define MIN_KEYPOINTS 5

#define keys \
    "{help h usage ? |        | print this message       }"\
    "{video v        | <none> | video to process         }"\
    "{@save s        |        | save data to disk or not }"\
    "{@load l        |        | load from disk or not    }"\
    "{count c        |    0   | number of images to load }"\
    "{sigma          |  0.03  | parameter: sigma         }"\
    "{scale          |  0.02  | parameter: scale         }"

void weightAdd(const cv::Mat &src1, const cv::Mat &src2, const cv::Mat &weight, cv::Mat &dst);

// VID_20160106_212608.mp4

int main(int argc, char **argv)
{
    cv::CommandLineParser parser(argc, argv, keys);
    if (parser.has("help"))
    {
        return 0;
    }

	std::vector<cv::Mat> images;
    std::vector<cv::Mat> Ms;

    Ms.push_back(cv::Mat());

    std::string video = parser.get<std::string>("video");
    bool load_from_disk = parser.has("@load");
    bool write_to_disk = parser.has("@save");
    int image_count = parser.get<int>("count");
    float sigma = parser.get<float>("sigma");
    float sigma2 = sigma * sigma;
    float scale = parser.get<float>("scale");

    std::cout << "sigma = " << sigma << std::endl;
    std::cout << "scale = " << scale << std::endl;

    if (load_from_disk)
    {
        std::cout << "Loading images from disk(image count:" << image_count << ")" << std::endl;

        images.push_back(cv::Mat());
        images[0] = cv::imread(IMAGE_FILE(0));
        if (images[0].empty())
        {
            std::cerr << "Cannot open " << IMAGE_FILE(0) << std::endl;
            return -1;
        }

        for (int idx = 1; idx < image_count; ++idx)
        {
            cv::Mat image;
            image = cv::imread(IMAGE_FILE(idx));
            if (image.empty())
                continue;

            cv::Mat M;
            cv::FileStorage fs(POSE_FILE(idx), cv::FileStorage::READ);
            fs["data"] >> M;

            images.push_back(image);
            Ms.push_back(M);
        }

    }
    else // if (load_from_disk)
    {

        cv::namedWindow(CAPTURE_WIN, cv::WINDOW_AUTOSIZE);
        cv::moveWindow(CAPTURE_WIN, 0, 0);

        cv::VideoCapture capture;
#ifdef __unix__
        capture.open(video);
#else
        capture.open(video);
#endif
        CV_Assert(capture.isOpened() == true);

        while (cv::waitKey(5) != 27)
        {
            cv::Mat src;
            capture >> src;
            if (src.empty())
            {
                std::cout << "End of capture" << std::endl;
                break;
            }

            cv::imshow(CAPTURE_WIN, src);
            images.push_back(src);
        }

        if (images.size() < 2)
            std::cout << "Too few images" << std::endl;
        else
            std::cout << images.size() << " image(s) totally." << std::endl;

        cv::destroyWindow(CAPTURE_WIN);

        cv::namedWindow(MATCH_WIN, cv::WINDOW_NORMAL);
        cv::resizeWindow(MATCH_WIN, images[0].size().width, images[0].size().height / 2);

        cv::Ptr<cv::xfeatures2d::SIFT> sift = cv::xfeatures2d::SIFT::create(100, 5);

        std::vector<cv::KeyPoint> keypoints_0;
        cv::Mat descriptors_0;
        sift->detectAndCompute(images[0], cv::Mat(), keypoints_0, descriptors_0);
        if (keypoints_0.size() < MIN_KEYPOINTS)
        {
            std::cout << "Not enough keypoints" << std::endl;
            return -1;
        }
        else
            std::cout << keypoints_0.size() << " keypoint(s) found" << std::endl;

        if( write_to_disk)
        cv::imwrite(IMAGE_FILE(0), images[0]);

        cv::FlannBasedMatcher matcher;

        std::vector<cv::KeyPoint> keypoints_idx;
        cv::Mat descriptors_idx;
        std::vector<cv::DMatch> matches;
        for (int idx = 1; idx < images.size(); ++idx)
        {
            keypoints_idx.clear();
            matches.clear();

            std::cout << "Handling " << idx << "/" << images.size() - 1 << std::endl;

            sift->detectAndCompute(images[idx], cv::Mat(), keypoints_idx, descriptors_idx);
            if (keypoints_idx.size() >= MIN_KEYPOINTS)
            {
                std::cout << keypoints_idx.size() << " keypoint(s) found" << std::endl;

                matcher.match(descriptors_0, descriptors_idx, matches);
                std::cout << matches.size() << " matche(s) found" << std::endl;
                if (matches.size() < 5)
                {
                    std::cout << "Not enough matches" << std::endl;
                    continue;
                }

                double max_dist = DBL_MIN;
                double min_dist = DBL_MAX;
                for (int i = 0; i < matches.size(); i++)
                {
                    double dist = matches[i].distance;
                    if (dist < min_dist) min_dist = dist;
                    if (dist > max_dist) max_dist = dist;
                }

                std::cout << "-- Max dist : " << max_dist << std::endl;
                std::cout << "-- Min dist : " << min_dist << std::endl;

                std::vector<cv::DMatch> good_matches;
#if 0
                for (int i = 0; i < matches.size(); i++)
                {
                    if (matches[i].distance < 4 * min_dist)
                    {
                        good_matches.push_back(matches[i]);
                    }
                }
#else
                for (int i = 0; i < 50; i++)
                    good_matches.push_back(matches[i]);
#endif

                cv::Mat match_canvas;
                cv::drawMatches(images[0], keypoints_0, images[idx], keypoints_idx, good_matches, match_canvas);
                cv::imshow(MATCH_WIN, match_canvas);

                std::vector<cv::Point2f> src_pts(good_matches.size());
                std::vector<cv::Point2f> dst_pts(good_matches.size());
                for (int i = 0; i < good_matches.size(); ++i)
                {
                    src_pts[i] = keypoints_0[good_matches[i].queryIdx].pt;
                    dst_pts[i] = keypoints_idx[good_matches[i].trainIdx].pt;
                }

                cv::Mat M = cv::findHomography(src_pts, dst_pts, CV_RANSAC);
                Ms.push_back(M);
                std::cout << M << std::endl;

                if (write_to_disk)
                {
                    cv::imwrite(IMAGE_FILE(idx), images[idx]);
                    cv::FileStorage fs(POSE_FILE(idx), cv::FileStorage::WRITE);
                    fs << "data" << M;
                }

            }
            else
            {
                std::cout << "Not enough keypoints" << std::endl;
            }

            //
            cv::waitKey(10);

        }

    } // LOAD_FROM_DISK
    
    CV_Assert(images.size() == Ms.size());
    std::cout << images.size() << " image(s) in total" << std::endl;

    cv::destroyAllWindows();

    cv::namedWindow(TOTAL_WIN, cv::WINDOW_AUTOSIZE);
    cv::moveWindow(TOTAL_WIN, 0, 0);
#ifdef DEBUG
        cv::namedWindow(BEFORE_WARPING_WIN, cv::WINDOW_NORMAL);
        cv::resizeWindow(BEFORE_WARPING_WIN, images[0].size().width / 2, images[0].size().height / 2);
        cv::moveWindow(BEFORE_WARPING_WIN, 0, images[0].size().height);

        cv::namedWindow(AFTER_WARPING_WIN, cv::WINDOW_NORMAL);
        cv::resizeWindow(AFTER_WARPING_WIN, images[0].size().width / 2, images[0].size().height / 2);
        cv::moveWindow(AFTER_WARPING_WIN, images[0].size().width / 2, images[0].size().height);

        cv::namedWindow(WEIGHT_WIN, cv::WINDOW_NORMAL);
        cv::resizeWindow(WEIGHT_WIN, images[0].size().width / 2, images[0].size().height / 2);
        cv::moveWindow(WEIGHT_WIN, images[0].size().width, 0);
#endif

    // Warp
    cv::Mat total(images[0].size(), CV_32FC3);
    total.setTo(0.0);
    for (int idx = 1; idx < images.size(); ++idx)
    {
        std::cout << "Merging " << idx << "/" << images.size() - 1 << std::endl;

        cv::Mat warpped;
#if 0
        cv::warpPerspective(images[idx], warpped, Ms[idx], images[0].size(), cv::WARP_INVERSE_MAP);
#else
        warpped = images[idx];
#endif

#ifdef DEBUG
        cv::imshow(BEFORE_WARPING_WIN, images[idx]);
        cv::imshow(AFTER_WARPING_WIN, warpped);
#endif

        cv::Mat mean;

#if 1
        cv::Mat hsv;
        cv::cvtColor(warpped, hsv, CV_BGR2HSV);
        std::vector<cv::Mat> hsvs;
        cv::split(hsv, hsvs);
        cv::Mat V;
        hsvs[2].convertTo(V, CV_32FC3);
        cv::Mat V_normalized;
        V_normalized = V / 255;

        cv::Mat warpped_32f;
        warpped.convertTo(warpped_32f, CV_32FC3);

        cv::Mat weight(total.size(), CV_32FC1);
        weight.setTo(1.0 / images.size());
        cv::Mat V_normalized_meaned = V_normalized - 1.0;
        cv::Mat V_normalized_meaned_2 = V_normalized_meaned.mul(V_normalized_meaned);
        cv::Mat V_normalized_meaned_2_exp;

        cv::exp(-V_normalized_meaned_2 / sigma2, V_normalized_meaned_2_exp);
        cv::add(weight, scale * V_normalized_meaned_2_exp, weight, cv::Mat(), CV_32FC1);

        cv::imshow(WEIGHT_WIN, weight);

        weightAdd(total, warpped_32f, weight, total);

        mean = total;
#else
        cv::add(total, dst, total, cv::Mat(), CV_32FC3);
        cv::Mat mean = total / idx;
#endif

        // Display
        cv::Mat canvas;
        mean.convertTo(canvas, CV_8UC3);
        cv::imshow(TOTAL_WIN, canvas);
        cv::waitKey(5);

        if (idx == images.size() - 1)
        {
            cv::imwrite("weight_" + std::to_string(sigma) + "_" + std::to_string(scale) + ".jpg", weight);
        }
    }

    // End
	cv::waitKey();

    // Save
    cv::Mat result;
    total.convertTo(result, CV_8UC3);
    cv::imwrite("final_" + std::to_string(sigma) + "_" + std::to_string(scale) + ".jpg", result);

	return 0;
}

void weightAdd(const cv::Mat &src1, const cv::Mat &src2, const cv::Mat &weight, cv::Mat &dst)
{
    CV_Assert (src1.size() == src2.size());
    CV_Assert (src1.type() == CV_32FC3 && src2.type() == CV_32FC3 && weight.type() == CV_32FC1);

    int rows = src1.size().height;
    int cols = src1.size().width;
    if (src1.isContinuous() && src2.isContinuous())
    {
        cols *= rows;
        rows = 1;
    }

    cv::Mat _dst(src1.size(), CV_32FC3);
    for (int r = 0; r < rows; ++r)
    {
        cv::Vec3f *src1_ptr   = (cv::Vec3f *)src1.  ptr(r);
        cv::Vec3f *src2_ptr   = (cv::Vec3f *)src2.  ptr(r);
        cv::Vec3f *dst_ptr    = (cv::Vec3f *)_dst.  ptr(r);
        float     *weight_ptr = (float *)    weight.ptr(r);

        for (int c = 0; c < cols; ++c, ++src1_ptr, ++src2_ptr, ++dst_ptr, ++weight_ptr)
        {
            *dst_ptr = *src1_ptr + *src2_ptr * *weight_ptr;
        }

    }

    dst = _dst;
}
